#ifndef AD_HPP_INCLUDED
#define AD_HPP_INCLUDED

/*
 * ad.hpp
 * This file is part of dynad
 *
 * Copyright (C) 2014 - Sebastian Gesemann
 *
 * dynad is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * dynad is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with dynad. If not, see <http://www.gnu.org/licenses/>.
 */

/*
 * Automatic Differentiation
 * (forward propagation + expression templates + move semantics)
 */

#include <map>
#include <utility>
#include <type_traits>
#include <cmath>
#include "cow.hpp"

namespace ad_details {

namespace more {

struct deriv_base
{};


template <class T> struct deriv_expr : deriv_base
{};


struct deriv_container;

struct deriv_single : deriv_expr<deriv_single>
{
	typedef std::pair<int, double> const* const_iterator;
	typedef std::pair<int, double>* iterator;

	std::pair<int, double> entry;

	explicit deriv_single(int dkey, double dval = 1.0) : entry(dkey, dval)
	{}

	inline void export_to(double factor, deriv_container& dc) const;

	const_iterator begin() const { return &entry; }
	const_iterator end() const { return &entry + 1; }
	const_iterator cbegin() const { return &entry; }
	const_iterator cend() const { return &entry + 1; }
	iterator begin() { return &entry; }
	iterator end() { return &entry + 1; }
};

struct deriv_container : deriv_expr<deriv_container>
{
	typedef std::map<int, double>::const_iterator const_iterator;
	typedef std::map<int, double>::iterator iterator;

	cow<std::map<int, double> > derivs;

	deriv_container()
	{}

	deriv_container(deriv_container const& x) : derivs(x.derivs)
	{}

	deriv_container(deriv_container&& x) : derivs(std::move(x.derivs))
	{}

	void reset()
	{
		derivs = cow<std::map<int, double> >();
	}

	void add_from(double factor, deriv_container const& dc);
	void add_from(double factor, deriv_single const& dc);

	void export_to(double factor, deriv_container& dc) const
	{
		dc.add_from(factor, *this);
	}

	const_iterator begin() const { return derivs.read().begin(); }
	const_iterator end() const { return derivs.read().end(); }
	const_iterator cbegin() const { return derivs.read().begin(); }
	const_iterator cend() const { return derivs.read().end(); }
	iterator begin() { return derivs.write().begin(); }
	iterator end() { return derivs.write().end(); }

	void add(int key, double dval = 1)
	{
		derivs.write()[key] += dval;
	}

	double get(int key) const {
		auto&& m = derivs.read();
		auto i = m.find(key);
		if (i==m.end()) return 0;
		return i->second;
	}
};

void deriv_single::export_to(double factor, deriv_container& dc) const
{
	dc.add_from(factor, *this);
}

void deriv_container::add_from(double factor, deriv_single const& dc)
{
	this->derivs.write()[dc.entry.first] += factor * dc.entry.second;
}

void deriv_container::add_from(double factor, deriv_container const& dc)
{
	auto& me = this->derivs.write();
	auto& other = dc.derivs.read();
	auto iter = other.begin();
	auto end = other.end();
	while (iter != end) {
		me[iter->first] += factor * iter->second;
		++iter;
	}
}

struct zero_derivs : deriv_expr<zero_derivs>
{
	void export_to(double factor, deriv_container& dc) const
	{}

};

template <class Derivs>
struct scaled_derivs : deriv_expr<scaled_derivs<Derivs> >
{
	double scale;
	Derivs derivs;

	template <class U>
	scaled_derivs(double s, U&& u)
	    : scale(s), derivs(std::forward<U>(u))
	{}


	scaled_derivs(scaled_derivs const& x) : scale(x.scale), derivs(x.derivs)
	{}

	scaled_derivs(scaled_derivs&& x)
	    : scale(x.scale), derivs(std::move(x.derivs))
	{}


	void export_to(double factor, deriv_container& dc) const
	{
		derivs.export_to(factor * scale, dc);
	}
};

template <class Derivs1, class Derivs2>
struct mixed_derivs : deriv_expr<mixed_derivs<Derivs1, Derivs2> >
{
	double scale1;
	double scale2;
	Derivs1 derivs1;
	Derivs2 derivs2;

	template <class U, class V>
	mixed_derivs(double s1, double s2, U&& u, V&& v)
	    : scale1(s1),
	      scale2(s2),
	      derivs1(std::forward<U>(u)),
	      derivs2(std::forward<V>(v))
	{}


	mixed_derivs(mixed_derivs const& x)
	    : scale1(x.scale1),
	      scale2(x.scale2),
	      derivs1(x.derivs1),
	      derivs2(x.derivs2)
	{}

	mixed_derivs(mixed_derivs&& x)
	    : scale1(x.scale1),
	      scale2(x.scale2),
	      derivs1(std::move(x.derivs1)),
	      derivs2(std::move(x.derivs2))
	{}


	void export_to(double factor, deriv_container& dc) const
	{
		derivs1.export_to(factor * scale1, dc);
		derivs2.export_to(factor * scale2, dc);
	}
};

inline zero_derivs operator*(zero_derivs, double)
{
	return zero_derivs();
}

inline deriv_single operator*(deriv_single const& x, double by)
{
	return deriv_single(x.entry.first, by * x.entry.second);
}

inline scaled_derivs<deriv_container> operator*(deriv_container const& x,
                                                double by)
{
	return scaled_derivs<deriv_container>(by, x);
}

inline scaled_derivs<deriv_container> operator*(deriv_container&& x, double by)
{
	return scaled_derivs<deriv_container>(by, std::move(x));
}

template <class Derivs>
inline scaled_derivs<Derivs> operator*(scaled_derivs<Derivs> const& x,
                                       double by)
{
	return scaled_derivs<Derivs>(x.scale * by, x.derivs);
}

template <class Derivs>
inline scaled_derivs<Derivs> operator*(scaled_derivs<Derivs>&& x, double by)
{
	return scaled_derivs<Derivs>(x.scale * by, std::move(x.derivs));
}

template <class D1, class D2>
inline mixed_derivs<D1, D2> operator*(mixed_derivs<D1, D2> const& x, double by)
{
	return mixed_derivs<D1, D2>(by * x.scale1, by * x.scale2, x.derivs1,
	                            x.derivs2);
}

template <class D1, class D2>
inline mixed_derivs<D1, D2> operator*(mixed_derivs<D1, D2>&& x, double by)
{
	return mixed_derivs<D1, D2>(by * x.scale1, by * x.scale2,
	                            std::move(x.derivs1), std::move(x.derivs2));
}

inline zero_derivs operator+(zero_derivs, zero_derivs)
{
	return zero_derivs();
}

template <class T>
inline typename std::decay<T>::type operator+(zero_derivs, T&& derivs)
{
	return std::forward<T>(derivs);
}

template <class T>
inline typename std::decay<T>::type operator+(T&& derivs, zero_derivs)
{
	return std::forward<T>(derivs);
}

template <class T> T const& unscale(deriv_expr<T> const& x)
{
	return static_cast<T const&>(x);
}
template <class T> T&& unscale(deriv_expr<T>&& x)
{
	return static_cast<T&&>(x);
}
template <class T> T const& unscale(scaled_derivs<T> const& x)
{
	return x.derivs;
}
template <class T> T&& unscale(scaled_derivs<T>&& x)
{
	return std::move(x.derivs);
}

double get_scale(deriv_base const&)
{
	return 1.0;
}
template <class T> double get_scale(scaled_derivs<T> const& x)
{
	return x.scale;
}

template <class T, bool Try> struct unscaled_type_helper
{
	typedef typename std::decay<decltype(unscale(std::declval<T>()))>::type
	type;
};
template <class T> struct unscaled_type_helper<T, false>
{};

template <class T>
struct unscaled_type
    : unscaled_type_helper<T,
                           std::is_convertible<T&&, deriv_base const&>::value>
{};


template <class T, class U>
mixed_derivs<typename unscaled_type<T>::type, typename unscaled_type<U>::type>
operator+(T&& x, U&& y)
{
	return mixed_derivs<typename unscaled_type<T>::type,
	                    typename unscaled_type<U>::type>(
	    get_scale(x), get_scale(y), unscale(std::forward<T>(x)),
	    unscale(std::forward<U>(y)));
}

} // namespace more

inline double value_of(double x)
{
	return x;
}

inline more::zero_derivs derivs_of(double)
{
	return more::zero_derivs();
}

template <class Derivs = more::deriv_container> struct ad
{
	double value;
	Derivs derivs;

	Derivs const& const_derivs() const
	{
		return derivs;
	}

	ad(double v = 0) : value(v)
	{}


	ad(double v, Derivs const& d) : value(v), derivs(d)
	{}


	ad(double v, Derivs&& d) : value(v), derivs(std::move(d))
	{}


	ad(ad const& x) : value(x.value), derivs(x.derivs)
	{}

	ad(ad&& x) : value(x.value), derivs(std::move(x.derivs))
	{}


	template <class D2> ad(ad<D2> const& x) : value(x.value)
	{
		x.derivs.export_to(1.0, this->derivs);
	}

	ad& operator=(ad const& x)
	{
		value = x.value;
		derivs = x.derivs;
		return *this;
	}
	ad& operator=(ad&& x)
	{
		value = x.value;
		derivs = std::move(x.derivs);
		return *this;
	}

	template <class D2>
	typename std::enable_if<!std::is_same<Derivs, D2>::value&& std::is_same<
	                             Derivs, more::deriv_container>::value,
	                        ad&>::type
	operator=(ad<D2> const& x)
	{
		this->value = x.value;
		this->derivs.reset();
		x.derivs.export_to(1.0, this->derivs);
	}

	friend double value_of(ad const& x)
	{
		return x.value;
	}
	friend Derivs const& derivs_of(ad const& x)
	{
		return x.derivs;
	}
	friend Derivs derivs_of(ad const&& x)
	{
		return x.derivs;
	}
	friend Derivs derivs_of(ad&& x)
	{
		return std::move(x.derivs);
	}
};

template <class D>
inline ad<typename std::decay<D>::type> make_ad(double v, D&& d)
{
	return ad<typename std::decay<D>::type>(v, std::forward<D>(d));
}

inline double make_ad(double v, more::zero_derivs)
{
	return v;
}

template <class T> struct is_ad
{
	static const bool value = false;
};
template <class D> struct is_ad<ad<D> >
{
	static const bool value = true;
};

template <class T> struct is_ad_aoref
{
	static const bool value = is_ad<typename std::decay<T>::type>::value;
};

template <class T, class U> struct at_least_one_ad
{
	static const bool value = is_ad_aoref<T>::value || is_ad_aoref<U>::value;
};

#define DEFINE_AD_COMPARISON(op)                                        \
	template<class T, class U, class = typename std::enable_if<         \
	                           at_least_one_ad<T, U>::value>::type>     \
	inline bool operator op(T const& x, U const& y)                     \
	{                                                                   \
	    return value_of(x) op value_of(y);                              \
	}

DEFINE_AD_COMPARISON(<)
DEFINE_AD_COMPARISON(>)
DEFINE_AD_COMPARISON(<=)
DEFINE_AD_COMPARISON(>=)
DEFINE_AD_COMPARISON(==)
DEFINE_AD_COMPARISON(!=)
#undef DEFINE_AD_COMPARISON

template <class T, class U,
          class = typename std::enable_if<at_least_one_ad<T, U>::value>::type>
inline auto operator+(T&& x, U&& y)
    -> decltype(make_ad(0.0, derivs_of(x) + derivs_of(y)))
{
	return make_ad(value_of(x) + value_of(y),
	               derivs_of(std::forward<T>(x)) +
	                   derivs_of(std::forward<U>(y)));
}

template <class T, class U,
          class = typename std::enable_if<at_least_one_ad<T, U>::value>::type>
inline auto operator-(T&& x, U&& y)
    -> decltype(make_ad(0.0, derivs_of(x) + derivs_of(y) * 0.0))
{
	return make_ad(value_of(x) - value_of(y),
	               derivs_of(std::forward<T>(x)) +
	                   derivs_of(std::forward<U>(y)) * (-1.0));
}

template <class T, class U,
          class = typename std::enable_if<at_least_one_ad<T, U>::value>::type>
inline auto operator*(T&& x, U&& y)
    -> decltype(make_ad(0.0, derivs_of(x) * 0.0 + derivs_of(y) * 0.0))
{
	return make_ad(value_of(x) * value_of(y),
	               derivs_of(std::forward<T>(x)) * value_of(y) +
	               derivs_of(std::forward<U>(y)) * value_of(x));
}

template <class T, class U,
          class = typename std::enable_if<at_least_one_ad<T, U>::value>::type>
inline auto operator/(T&& x, U&& y)
    -> decltype(make_ad(0.0, derivs_of(x) * 0.0 + derivs_of(y) * 0.0))
{
	const double invy = 1.0 / value_of(y);
	const double divr = value_of(x) * invy;
	return make_ad(divr, derivs_of(std::forward<T>(x)) * invy +
	                     derivs_of(std::forward<U>(y)) * (-divr * invy));
}

template <class T, class = typename std::enable_if<is_ad_aoref<T>::value>::type>
inline auto sin(T&& x) -> decltype(make_ad(0.0, derivs_of(x) * 0.0))
{
	double s = std::sin(value_of(x));
	double c = std::cos(value_of(x));
	return make_ad(s, derivs_of(std::forward<T>(x)) * c);
}

template <class T, class = typename std::enable_if<is_ad_aoref<T>::value>::type>
inline auto cos(T&& x) -> decltype(make_ad(0.0, derivs_of(x) * 0.0))
{
	double s = std::sin(value_of(x));
	double c = std::cos(value_of(x));
	return make_ad(c, derivs_of(std::forward<T>(x)) * (-s));
}

template <class T, class = typename std::enable_if<is_ad_aoref<T>::value>::type>
inline auto square(T&& x) -> decltype(make_ad(0.0, derivs_of(x) * 0.0))
{
	double v = value_of(x);
	return make_ad(v * v, derivs_of(std::forward<T>(x)) * (2 * v));
}

template <class T, class = typename std::enable_if<is_ad_aoref<T>::value>::type>
inline auto sqrt(T&& x) -> decltype(make_ad(0.0, derivs_of(x) * 0.0))
{
	double v = std::sqrt(value_of(x));
	return make_ad(v, derivs_of(std::forward<T>(x)) * (0.5 / v));
}

template <class T, class = typename std::enable_if<is_ad_aoref<T>::value>::type>
inline auto exp(T&& x) -> decltype(make_ad(0.0, derivs_of(x) * 0.0))
{
	double v = std::exp(value_of(x));
	return make_ad(v, derivs_of(std::forward<T>(x)) * v);
}

template <class T, class = typename std::enable_if<is_ad_aoref<T>::value>::type>
inline auto log(T&& x) -> decltype(make_ad(0.0, derivs_of(x) * 0.0))
{
	double v = value_of(x);
	return make_ad(std::log(v), derivs_of(std::forward<T>(x)) * (1.0 / v));
}

template <class T, class = typename std::enable_if<is_ad_aoref<T>::value>::type>
inline auto abs(T&& x) -> decltype(make_ad(0.0, derivs_of(x) * 0.0))
{
	double v = value_of(x);
	double s = v >= 0 ? 1 : -1;
	return make_ad(v * s, derivs_of(std::forward<T>(x)) * s);
}

template <class D> ad<>& operator+=(ad<>& io, ad<D> const& x)
{
	io.value += x.value();
	x.derivs.export_to(1.0, io);
	return io;
}

template <class D> ad<>& operator-=(ad<>& io, ad<D> const& x)
{
	io.value -= x.value();
	x.derivs.export_to(-1.0, io);
	return io;
}

template <class D> ad<>& operator*=(ad<>& io, ad<D> const& x)
{
	io = std::move(io) * x;
	return io;
}

template <class D> ad<>& operator/=(ad<>& io, ad<D> const& x)
{
	io = std::move(io) / x;
	return io;
}

} // namespace ad_details

typedef ad_details::ad<> advar;

inline ad_details::ad<ad_details::more::deriv_single> adiff(double value,
                                                            int dkey)
{
	return ad_details::ad<ad_details::more::deriv_single>(
	    value, ad_details::more::deriv_single(dkey));
}

#endif

